/**
* \brief 
* \author pengcheng (pengcheng@yslrpch@126.com)
* \date 2020-05-30
* \attention CopyrightÃ‚Â©ADC Technology(tianjin)Co.Ltd
* \attention Refer to COPYRIGHT.txt for complete terms of copyright notice
*/

#include "detection_vision/tensorrt/calibration.h"

namespace tensorrt_inference
{
Int8CacheCalibrator::Int8CacheCalibrator(BatchStream& stream,
                                         int firstBatch,
                                         std::string calibrationTableName,
                                         bool readCache):
    stream_(stream),
    calibration_table_name_(calibrationTableName),
    read_cache_(readCache)
{
    if (calibration_table_name_.empty())
    {
        return;
    }
    LOG(INFO) << "create calibrator with calibration table: " << calibration_table_name_;
    nvinfer1::DimsNCHW dims = stream_.getDims();
    input_count_ = stream_.getBatchSize() * dims.c() * dims.h() * dims.w();
    CUDACHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float)));
    stream_.reset(firstBatch);
}

int Int8CacheCalibrator::getBatchSize() const
{
    return stream_.getBatchSize();
}

bool Int8CacheCalibrator::getBatch(void* bindings[], const char* names[], int nbBindings)
{
    (void) nbBindings;
    LOG(INFO) << "get batch for input blob: " << names[0];
    if (!stream_.next())
    {
        return false;
    }
    CUDACHECK(cudaMemcpy(device_input_, stream_.getBatch(), input_count_ * sizeof(float), cudaMemcpyHostToDevice));
    bindings[0] = device_input_;
    return true;
}

const void* Int8CacheCalibrator::readCalibrationCache(size_t& length)
{
    LOG(INFO) << "read calibration cache from calibration table: " << calibration_table_name_;
    calibration_cache_.clear();
    std::ifstream input(calibration_table_name_, std::ios::binary);
    input >> std::noskipws;
    if (read_cache_ && input.good())
    {
        std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(), std::back_inserter(calibration_cache_));
    }
    length = calibration_cache_.size();
    return length ? calibration_cache_.data() : nullptr;
}

void Int8CacheCalibrator::writeCalibrationCache(const void* cache, size_t length)
{
    LOG(INFO) << "write calibration cache into calibration table: " << calibration_table_name_;
    std::ofstream output(calibration_table_name_, std::ios::binary);
    if (!output.is_open())
    {
        std::string error = calibration_table_name_ + " open file failed";
        LOG(ERROR) << error;
        throw error;
    }
    output.write(reinterpret_cast<const char*>(cache), length);
    output.close();
}

}  